package com.test.service;

import com.test.Utils.StockTradeMapper;
import com.test.entity.StockTrade;
import lombok.extern.slf4j.Slf4j;
import org.apache.flink.api.common.eventtime.WatermarkStrategy;
import org.apache.flink.api.common.functions.AggregateFunction;
import org.apache.flink.api.common.functions.MapFunction;
import org.apache.flink.api.common.serialization.SimpleStringSchema;
import org.apache.flink.api.java.tuple.Tuple2;
import org.apache.flink.configuration.Configuration;
import org.apache.flink.connector.kafka.source.KafkaSource;
import org.apache.flink.connector.kafka.source.enumerator.initializer.OffsetsInitializer;
import org.apache.flink.streaming.api.datastream.DataStream;
import org.apache.flink.streaming.api.environment.StreamExecutionEnvironment;
import org.apache.flink.streaming.api.functions.sink.RichSinkFunction;
import org.apache.flink.streaming.api.windowing.assigners.TumblingEventTimeWindows;
import org.apache.flink.streaming.api.windowing.time.Time;
import org.apache.flink.streaming.connectors.kafka.FlinkKafkaConsumer;
import org.apache.flink.streaming.connectors.redis.common.config.FlinkJedisPoolConfig;
import redis.clients.jedis.Jedis;

import java.time.Duration;
import java.util.Properties;

@Slf4j
public class TopTenPlatform {
    public static FlinkJedisPoolConfig redisConfig = new FlinkJedisPoolConfig.Builder()
            .setHost("192.168.43.150")
            .setPort(6379)
            .build();

    public static void createDataStream(StreamExecutionEnvironment env) {
        Properties properties = new Properties();
        properties.setProperty("bootstrap.servers", "192.168.88.135:19092,192.168.88.135:29092,192.168.88.135:39092");
        properties.setProperty("group.id", "top-ten-platform");

        FlinkKafkaConsumer<String> consumer = new FlinkKafkaConsumer<>(
            "stock-trades-",
            new SimpleStringSchema(),
            properties
        );
        consumer.setStartFromEarliest();

        // 创建数据流并添加水印
        DataStream<StockTrade> tradeStream = env
                .addSource(consumer)
                .map(new StockTradeMapper())
                .assignTimestampsAndWatermarks(WatermarkStrategy
                        .<StockTrade>forBoundedOutOfOrderness(Duration.ofSeconds(1))
                        .withTimestampAssigner((event, timestamp) -> event.getTime())
                );

        // 当天
        tradeStream
                .map(new MapFunction<StockTrade, Tuple2<String, Double>>() {
                    @Override
                    public Tuple2<String, Double> map(StockTrade value) throws Exception {
                        return new Tuple2<>(value.getTradePlatform(), value.getTradeVolume()*value.getPrice());
                    }
                })
                .keyBy(value -> value.f0)
                .window(TumblingEventTimeWindows.of(Time.seconds(1)))
                .aggregate(new TradeAggregator())
                .addSink(new AccumulatingRedisSink(redisConfig));

        log.info("Top ten platform analysis configured");
    }

    // 自定义 Sink 函数
    public static class AccumulatingRedisSink extends RichSinkFunction<Tuple2<String, Double>> {
        private transient Jedis jedis;
        private final FlinkJedisPoolConfig config;

        public AccumulatingRedisSink(FlinkJedisPoolConfig config) {
            this.config = config;
        }

        @Override
        public void open(Configuration parameters) throws Exception {
            super.open(parameters);
            jedis = new Jedis(config.getHost(), config.getPort());
        }

        @Override
        public void close() throws Exception {
            if (jedis != null) {
                jedis.close();
            }
            super.close();
        }

        @Override
        public void invoke(Tuple2<String, Double> value, Context context) throws Exception {

            String hashKey = "stock_trade_platform"; // 存储股票交易量的键
            String field = value.f0; // 使用股票平台作为字段

            // 检查键的类型
            String type = jedis.type(hashKey);
            if (!"zset".equals(type)) {
                // 如果键不存在或类型不正确，删除旧的键
                jedis.del(hashKey);
                // 创建新的有序集合
                jedis.zadd(hashKey, 0, field); // 初始化一个空的有序集合
            }

            // 从 Redis 中获取现有值
            Double existingVolume = jedis.zscore(hashKey, field);
            if (existingVolume == null) {
                existingVolume = (double) 0L;
            }

            // 更新交易量
            double updatedVolume = existingVolume + value.f1;

            // 将更新后的值写回 Redis
            jedis.zadd(hashKey, updatedVolume, field);
        }

    }

    // 定义聚合函数
    public static class TradeAggregator implements AggregateFunction<Tuple2<String, Double>, Tuple2<String, Double>, Tuple2<String, Double>> {

        @Override
        public Tuple2<String, Double> createAccumulator() {
            return new Tuple2<>("", 0d);
        }

        @Override
        public Tuple2<String, Double> add(Tuple2<String, Double> stringDoubleTuple2, Tuple2<String, Double> stringDoubleTuple22) {
            String stockCode = stringDoubleTuple2.f0;
            Double totalAmount = stringDoubleTuple2.f1 + stringDoubleTuple22.f1;
            return new Tuple2<>(stockCode, totalAmount);
        }

        @Override
        public Tuple2<String, Double> getResult(Tuple2<String, Double> accumulator) {
            return accumulator;
        }

        @Override
        public Tuple2<String, Double> merge(Tuple2<String, Double> a, Tuple2<String, Double> b) {
            String stockCode = a.f0; // 假设两个 accumulator 的股票代码相同
            Double totalAmount = a.f1 + b.f1;
            return new Tuple2<>(stockCode, totalAmount);
        }

    }
}
